package com.wh.springbootssm.utils;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.AmbiguousTableNameException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlProducer;
import org.dbunit.operation.DatabaseOperation;
import org.xml.sax.InputSource;

import java.io.*;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * @Auther: HuiEr
 * @Date: 2020/2/21 19:35
 * @Description
 */
public class AbstractDbunitTestCase {

    //需要从外部传入，里面多个方法都要调用
    private DatabaseConnection conn;
    //不需要从外部传入，里面多个方法需要调用
    private File tempFile;
    //测试数据集，需要从外部传入
    private IDataSet dataSet;

    /*public AbstractDbunitTestCase(){

    }*/

    public AbstractDbunitTestCase(String testFilePath) throws DataSetException {
        this.dataSet=new FlatXmlDataSet(new FlatXmlProducer(new InputSource(AbstractDbunitTestCase.class.getClassLoader().getResourceAsStream(testFilePath))));
    }
    public void setDatabaseConn(Connection conn) throws DatabaseUnitException {
        this.conn=new DatabaseConnection(conn);
    }

    public void backManyTables(String... tabNames) throws DataSetException, IOException {
        QueryDataSet queryDataSet = new QueryDataSet(conn);

        for (String tabName:tabNames) {
            queryDataSet.addTable(tabName);
        }

        tempFile = File.createTempFile("back", ".xml");

        FlatXmlDataSet.write(queryDataSet,new FileOutputStream(tempFile));
    }

    public void backOneTable(String tabName) throws IOException, DataSetException {
        backManyTables(tabName);
    }

   /* public void backTablesAsTest(String tabName) throws DataSetException, IOException {
        QueryDataSet queryDataSet = new QueryDataSet(conn);
        queryDataSet.addTable("user");
        File file111 = new File("test.xml");

        FlatXmlDataSet.write(queryDataSet,new FileOutputStream(file111));
    }*/
    public void insertTestData() throws DatabaseUnitException, SQLException {
        DatabaseOperation.CLEAN_INSERT.execute(conn,dataSet);
    }

    public void resumeTable() throws FileNotFoundException, DatabaseUnitException, SQLException {
        FlatXmlDataSet tmpDataSet = new FlatXmlDataSet(new FlatXmlProducer(new InputSource(new FileInputStream(tempFile))));
        DatabaseOperation.CLEAN_INSERT.execute(conn,tmpDataSet);
    }

}
